package org.example.jimeng;


import com.alibaba.fastjson.JSONObject;
import org.apache.commons.codec.binary.Hex;
import com.google.common.io.ByteStreams;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.ByteBuffer;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.security.MessageDigest;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.Base64;

/**
 * Copyright (year) Beijing Volcano Engine Technology Ltd.
 * <p>
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
public class Sign {
    private static final BitSet URLENCODER = new BitSet(256);
    private static final String CONST_ENCODE = "0123456789ABCDEF";
    public static final Charset UTF_8 = StandardCharsets.UTF_8;
    private final String region;
    private final String service;
    private final String schema;
    private final String host;
    private final String path;
    private final String ak;
    private final String sk;

    static {
        int i;
        for (i = 97; i <= 122; ++i) {
            URLENCODER.set(i);
        }
        for (i = 65; i <= 90; ++i) {
            URLENCODER.set(i);
        }
        for (i = 48; i <= 57; ++i) {
            URLENCODER.set(i);
        }
        URLENCODER.set('-');
        URLENCODER.set('_');
        URLENCODER.set('.');
        URLENCODER.set('~');
    }

    public static void main(String[] args) throws Exception {
        // 火山官网密钥信息, 注意sk结尾有==
        String AccessKeyID = "AKLTMmExYmRkNTI4NjA5NDY2ZTkzZjFiMTkzNTA5YjkyNzg";
        String SecretAccessKey = "WlRaak56RmxOalk0WXpZeE5HWmpNamhsTmpkbE5qVXdOakpoWlRkaVlUaw==";
        // 请求域名
        String endpoint = "visual.volcengineapi.com";
        String path = "/"; // 路径，不包含 Query// 请求接口信息
        String service = "cv";
        String region = "cn-north-1";
        String schema = "https";
        Sign sign = new Sign(region, service, schema, endpoint, path, AccessKeyID, SecretAccessKey);
        
        // 测试图片生成
        String prompt = "一只可爱的小猫";
        String result = sign.generateImage(prompt);
        System.out.println("生成结果: " + result);
        
        // 测试图片处理（如果有输入图片）
        // String inputImagePath = "src/main/resources/image/input.jpg";
        // String result2 = sign.processImage(inputImagePath, "将图片转换为卡通风格");
        // System.out.println("处理结果: " + result2);
    }

    public Sign(String region, String service, String schema, String host, String path, String ak, String sk) {
        this.region = region;
        this.service = service;
        this.host = host;
        this.schema = schema;
        this.path = path;
        this.ak = ak;
        this.sk = sk;
    }

    public void doRequest(String method, Map<String, String> queryList, byte[] body,
                          Date date, String action, String version) throws Exception {
        if (body == null) {
            body = new byte[0];
        }
        String xContentSha256 = hashSHA256(body);
        SimpleDateFormat sdf = new SimpleDateFormat("yyyyMMdd'T'HHmmss'Z'");
        sdf.setTimeZone(TimeZone.getTimeZone("GMT"));
        String xDate = sdf.format(date);
        String shortXDate = xDate.substring(0, 8);
        String contentType = "application/json";
        String signHeader = "host;x-date;x-content-sha256;content-type";

        SortedMap<String, String> realQueryList = new TreeMap<>(queryList);
        realQueryList.put("Action", action);
        realQueryList.put("Version", version);
        StringBuilder querySB = new StringBuilder();
        for (String key : realQueryList.keySet()) {
            querySB.append(signStringEncoder(key)).append("=").append(signStringEncoder(realQueryList.get(key))).append("&");
        }
        querySB.deleteCharAt(querySB.length() - 1);

        String canonicalStringBuilder = method + "\n" + path + "\n" + querySB + "\n" +
                "host:" + host + "\n" +
                "x-date:" + xDate + "\n" +
                "x-content-sha256:" + xContentSha256 + "\n" +
                "content-type:" + contentType + "\n" +
                "\n" +
                signHeader + "\n" +
                xContentSha256;

        System.out.println(canonicalStringBuilder);

        String hashcanonicalString = hashSHA256(canonicalStringBuilder.getBytes());
        String credentialScope = shortXDate + "/" + region + "/" + service + "/request";
        String signString = "HMAC-SHA256" + "\n" + xDate + "\n" + credentialScope + "\n" + hashcanonicalString;

        byte[] signKey = genSigningSecretKeyV4(sk, shortXDate, region, service);
//        String signature = HexFormat.of().formatHex(hmacSHA256(signKey, signString));
        String signature = Hex.encodeHexString(hmacSHA256(signKey, signString));

        URL url = new URL(schema + "://" + host + path + "?" + querySB);

        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod(method);
        conn.setRequestProperty("Host", host);
        conn.setRequestProperty("X-Date", xDate);
        conn.setRequestProperty("X-Content-Sha256", xContentSha256);
        conn.setRequestProperty("Content-Type", contentType);
        conn.setRequestProperty("Authorization", "HMAC-SHA256" +
                " Credential=" + ak + "/" + credentialScope +
                ", SignedHeaders=" + signHeader +
                ", Signature=" + signature);
        if (!Objects.equals(conn.getRequestMethod(), "GET")) {
            conn.setDoOutput(true);
            OutputStream os = conn.getOutputStream();
            os.write(body);
            os.flush();
            os.close();
        }
        conn.connect();

        int responseCode = conn.getResponseCode();

        InputStream is;
        if (responseCode == 200) {
            is = conn.getInputStream();
        } else {
            is = conn.getErrorStream();
        }
        String responseBody = new String(ByteStreams.toByteArray(is));
        is.close();

        System.out.println(responseCode);
        System.out.println(responseBody);
    }

    private String signStringEncoder(String source) {
        if (source == null) {
            return null;
        }
        StringBuilder buf = new StringBuilder(source.length());
        ByteBuffer bb = UTF_8.encode(source);
        while (bb.hasRemaining()) {
            int b = bb.get() & 255;
            if (URLENCODER.get(b)) {
                buf.append((char) b);
            } else if (b == 32) {
                buf.append("%20");
            } else {
                buf.append("%");
                char hex1 = CONST_ENCODE.charAt(b >> 4);
                char hex2 = CONST_ENCODE.charAt(b & 15);
                buf.append(hex1);
                buf.append(hex2);
            }
        }

        return buf.toString();
    }

    public static String hashSHA256(byte[] content) throws Exception {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");


            return Hex.encodeHexString(md.digest(content));
        } catch (Exception e) {
            throw new Exception(
                    "Unable to compute hash while signing request: "
                            + e.getMessage(), e);
        }
    }

    public static byte[] hmacSHA256(byte[] key, String content) throws Exception {
        try {
            Mac mac = Mac.getInstance("HmacSHA256");
            mac.init(new SecretKeySpec(key, "HmacSHA256"));
            return mac.doFinal(content.getBytes());
        } catch (Exception e) {
            throw new Exception(
                    "Unable to calculate a request signature: "
                            + e.getMessage(), e);
        }
    }

    private byte[] genSigningSecretKeyV4(String secretKey, String date, String region, String service) throws Exception {
        byte[] kDate = hmacSHA256((secretKey).getBytes(), date);
        byte[] kRegion = hmacSHA256(kDate, region);
        byte[] kService = hmacSHA256(kRegion, service);
        return hmacSHA256(kService, "request");
    }
    
    /**
     * 生成图片的便捷方法
     * @param prompt 图片描述
     * @return 生成的图片路径
     */
    public String generateImage(String prompt) throws Exception {
        String action = "CVProcess";
        String version = "2022-08-31";
        Date date = new Date();
        
        JSONObject req = new JSONObject();
        req.put("req_key", "jimeng_high_aes_general_v21_L");
        req.put("prompt", prompt);
        // 不添加image_urls，表示纯文本生成图片
        
        String response = doRequestWithResponse("POST", new HashMap<>(), req.toJSONString().getBytes(), date, action, version);
        
        // 解析响应并保存图片
        return saveGeneratedImage(response, prompt);
    }
    
    /**
     * 处理现有图片的便捷方法
     * @param inputImagePath 输入图片路径
     * @param prompt 处理描述
     * @return 处理后的图片路径
     */
    public String processImage(String inputImagePath, String prompt) throws Exception {
        String action = "CVProcess";
        String version = "2022-08-31";
        Date date = new Date();
        
        // 读取并编码输入图片
        String base64Image = encodeImageToBase64(inputImagePath);
        
        JSONObject req = new JSONObject();
        req.put("req_key", "jimeng_high_aes_general_v21_L");
        ArrayList<String> imageUrls = new ArrayList<>();
        imageUrls.add("data:image/jpeg;base64," + base64Image);
        req.put("image_urls", imageUrls);
        req.put("prompt", prompt);
        
        String response = doRequestWithResponse("POST", new HashMap<>(), req.toJSONString().getBytes(), date, action, version);
        
        // 解析响应并保存图片
        return saveGeneratedImage(response, prompt);
    }
    
    /**
     * 带响应返回的请求方法
     */
    private String doRequestWithResponse(String method, Map<String, String> queryList, byte[] body,
                                       Date date, String action, String version) throws Exception {
        if (body == null) {
            body = new byte[0];
        }
        String xContentSha256 = hashSHA256(body);
        SimpleDateFormat sdf = new SimpleDateFormat("yyyyMMdd'T'HHmmss'Z'");
        sdf.setTimeZone(TimeZone.getTimeZone("GMT"));
        String xDate = sdf.format(date);
        String shortXDate = xDate.substring(0, 8);
        String contentType = "application/json";
        String signHeader = "host;x-date;x-content-sha256;content-type";

        SortedMap<String, String> realQueryList = new TreeMap<>(queryList);
        realQueryList.put("Action", action);
        realQueryList.put("Version", version);
        StringBuilder querySB = new StringBuilder();
        for (String key : realQueryList.keySet()) {
            querySB.append(signStringEncoder(key)).append("=").append(signStringEncoder(realQueryList.get(key))).append("&");
        }
        querySB.deleteCharAt(querySB.length() - 1);

        String canonicalStringBuilder = method + "\n" + path + "\n" + querySB + "\n" +
                "host:" + host + "\n" +
                "x-date:" + xDate + "\n" +
                "x-content-sha256:" + xContentSha256 + "\n" +
                "content-type:" + contentType + "\n" +
                "\n" +
                signHeader + "\n" +
                xContentSha256;

        String hashcanonicalString = hashSHA256(canonicalStringBuilder.getBytes());
        String credentialScope = shortXDate + "/" + region + "/" + service + "/request";
        String signString = "HMAC-SHA256" + "\n" + xDate + "\n" + credentialScope + "\n" + hashcanonicalString;

        byte[] signKey = genSigningSecretKeyV4(sk, shortXDate, region, service);
        String signature = Hex.encodeHexString(hmacSHA256(signKey, signString));

        URL url = new URL(schema + "://" + host + path + "?" + querySB);

        HttpURLConnection conn = (HttpURLConnection) url.openConnection();
        conn.setRequestMethod(method);
        conn.setRequestProperty("Host", host);
        conn.setRequestProperty("X-Date", xDate);
        conn.setRequestProperty("X-Content-Sha256", xContentSha256);
        conn.setRequestProperty("Content-Type", contentType);
        conn.setRequestProperty("Authorization", "HMAC-SHA256" +
                " Credential=" + ak + "/" + credentialScope +
                ", SignedHeaders=" + signHeader +
                ", Signature=" + signature);
        if (!Objects.equals(conn.getRequestMethod(), "GET")) {
            conn.setDoOutput(true);
            OutputStream os = conn.getOutputStream();
            os.write(body);
            os.flush();
            os.close();
        }
        conn.connect();

        int responseCode = conn.getResponseCode();

        InputStream is;
        if (responseCode == 200) {
            is = conn.getInputStream();
        } else {
            is = conn.getErrorStream();
        }
        String responseBody = new String(ByteStreams.toByteArray(is));
        is.close();

        System.out.println("响应状态码: " + responseCode);
        System.out.println("响应内容: " + responseBody);
        
        return responseBody;
    }
    
    /**
     * 将图片编码为Base64
     */
    private String encodeImageToBase64(String imagePath) throws IOException {
        Path path = Paths.get(imagePath);
        byte[] imageBytes = Files.readAllBytes(path);
        return Base64.getEncoder().encodeToString(imageBytes);
    }
    
    /**
     * 保存生成的图片
     */
    private String saveGeneratedImage(String response, String prompt) throws Exception {
        // 解析响应JSON
        JSONObject responseJson = JSONObject.parseObject(response);
        
        // 检查响应状态
        if (responseJson.containsKey("code") && responseJson.getInteger("code") != 200) {
            throw new Exception("API请求失败: " + responseJson.getString("message"));
        }
        
        // 获取图片数据
        JSONObject data = responseJson.getJSONObject("data");
        if (data == null) {
            throw new Exception("响应中没有图片数据");
        }
        
        // 获取图片URL或Base64数据
        String imageUrl = data.getString("image_url");
        String imageBase64 = data.getString("image_base64");
        
        // 确保image文件夹存在
        String imageDir = "src/main/resources/image";
        Path imagePath = Paths.get(imageDir);
        if (!Files.exists(imagePath)) {
            Files.createDirectories(imagePath);
        }
        
        // 生成文件名
        String timestamp = new SimpleDateFormat("yyyyMMdd_HHmmss").format(new Date());
        String safePrompt = prompt.replaceAll("[^a-zA-Z0-9\\u4e00-\\u9fa5]", "_");
        String fileName = "generated_" + timestamp + "_" + safePrompt.substring(0, Math.min(safePrompt.length(), 20)) + ".jpg";
        String filePath = imageDir + "/" + fileName;
        
        // 保存图片
        if (imageBase64 != null && !imageBase64.isEmpty()) {
            // 保存Base64图片
            byte[] imageBytes = Base64.getDecoder().decode(imageBase64);
            Files.write(Paths.get(filePath), imageBytes);
        } else if (imageUrl != null && !imageUrl.isEmpty()) {
            // 下载图片URL
            downloadImage(imageUrl, filePath);
        } else {
            throw new Exception("响应中没有有效的图片数据");
        }
        
        System.out.println("图片已保存到: " + filePath);
        return filePath;
    }
    
    /**
     * 下载图片
     */
    private void downloadImage(String imageUrl, String filePath) throws IOException {
        URL url = new URL(imageUrl);
        try (InputStream in = url.openStream();
             FileOutputStream out = new FileOutputStream(filePath)) {
            byte[] buffer = new byte[1024];
            int bytesRead;
            while ((bytesRead = in.read(buffer)) != -1) {
                out.write(buffer, 0, bytesRead);
            }
        }
    }
}
